#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：academic_trend_analysis
@File    ：llm_analyzer.py
@IDE     ：PyCharm
@Author  ：iyoahs
@Date    ：2025/6/20 10:45
@Describe：LLM分析模块
'''
from typing import List, Dict
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.schema import StrOutputParser
from src.utils.logger import setup_logger

logger = setup_logger()


class LLMAnalyzer:
    """LLM分析器"""

    def __init__(self, config: Dict):
        self.llm = ChatOpenAI(
            model=config["model_name"],
            api_key=config["api_key"],
            base_url=config["base_url"],
            max_tokens=config["max_tokens"],
        )

    def generate_summary(self, title: str, abstract: str) -> str:
        """生成论文摘要"""
        prompt = PromptTemplate.from_template("""
请为以下学术论文生成简洁的摘要(100字以内):

标题: {title}
摘要: {abstract}

摘要要求:
- 突出核心贡献
- 使用学术语言
- 不超过100字
- 使用中文
        """)

        chain = prompt | self.llm | StrOutputParser()
        return chain.invoke({
            "title": title,
            "abstract": abstract
        })

    def generate_trend_report(self, keywords: List[Dict]) -> str:
        """生成趋势报告"""
        prompt = PromptTemplate.from_template("""
根据以下关键词及其出现频率，生成AI研究趋势报告:

{keywords}

报告要求:
1. 当前研究热点概述 (100字)
2. 主要研究主题分析 (200字)
3. 未来一年新兴趋势预测 (100字)
4. 使用中文输出
        """)

        keyword_str = "\n".join([f"{k['keyword']}: {k['count']}篇" for k in keywords])
        chain = prompt | self.llm | StrOutputParser()
        return chain.invoke({"keywords": keyword_str})

    def generate_paper_reports(self, papers: List[Dict]) -> List[Dict]:
        """为多篇论文生成报告"""
        return [{
            "id": p["id"],
            "summary": self.generate_summary([p])
        } for p in papers]

    def generate_report(self, topic: str, papers: List[Dict]) -> str:
        """根据多篇相同主题的论文写一个近期研究综述"""
        papers_abstract = ""
        for idx, paper in enumerate(papers):
            papers_abstract += (f"### 论文{idx+1}:" + paper["title"].replace("\n", " ")) + "\n"
            papers_abstract += paper["abstract"].replace("\n", " ")
            papers_abstract += "\n"
        prompt = PromptTemplate.from_template("""
下面，我会为你提供多篇关于以"{topic}"为主题的论文摘要，这些都是最近的研究成果，请你根据这些研究成果，生成一个近期研究报告，下面是这些文章的摘要:

{papers_abstract}

报告要求：
1. 当前研究热点概述 (100字)
2. 主要研究主题分析 (100字)
3. 该研究领域研究趋势的预测(200字)
4. 提供一些本研究主题的一些可行的研究方向(100字)
5. 输出的结果必须是中文，以markdown的格式排版输出
        """)

        chain = prompt | self.llm | StrOutputParser()
        return chain.invoke({"topic": topic, "papers_abstract": papers_abstract})